/*
 * Copyright © 2022-now 尛飛俠（Denvie） All rights reserved.
 */

package cn.denvie.springboot.dds.config;

import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Primary;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 数据源路由配置。继承 AbstractRoutingDataSource 提供两个库 “ds1”、“ds2”。
 *
 * @author Denvie
 * @date 2022/3/28
 * @since 1.0.0
 */
@Primary
@Component
@ConditionalOnProperty(name = "spring.profiles.active", havingValue = "jdbc")
public class MyRoutingDataSource extends AbstractRoutingDataSource {
    /**
     * 当前线程数据源名称
     */
    private final ThreadLocal<String> dsName = new ThreadLocal<>();

    @Resource(name = "dataSource1")
    private DataSource dataSource1;
    @Resource(name = "dataSource2")
    private DataSource dataSource2;

    @Override
    protected Object determineCurrentLookupKey() {
        return this.dsName.get();
    }

    /**
     * AbstractRoutingDataSource 实现了 InitializingBean 接口，afterPropertiesSet() 在属性初始化完成后调用。
     */
    @Override
    public void afterPropertiesSet() {
        // 初始化所有的目标数据源
        Map<Object, Object> targetDataSource = new HashMap<>(2);
        targetDataSource.put("ds1", dataSource1);
        targetDataSource.put("ds2", dataSource2);
        super.setTargetDataSources(targetDataSource);

        // 设置默认数据源
        super.setDefaultTargetDataSource(dataSource1);

        super.afterPropertiesSet();
    }

    public ThreadLocal<String> getDsName() {
        return dsName;
    }

    @PreDestroy
    public void preDestroy() {
        // 必须回收自定义的ThreadLocal变量
        this.dsName.remove();
    }
}
